package cn.iocoder.yudao.module.chat.config.gpt;

import cn.hutool.core.util.StrUtil;
import com.unfbx.chatgpt.OpenAiStreamClient;
import com.unfbx.chatgpt.function.KeyRandomStrategy;
import com.unfbx.chatgpt.interceptor.DynamicKeyOpenAiAuthInterceptor;
import com.unfbx.chatgpt.interceptor.OpenAILogger;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.annotation.Resource;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.concurrent.TimeUnit;

@Configuration
@EnableConfigurationProperties(ChatGPTProperties.class)
public class ChatGPTClientConfiguration {
    @Resource
    private ChatGPTProperties chatGPTProperties;
    @Bean
    public OpenAiStreamClient openAiStreamClient() {
        //本地开发需要配置代理地址
        Proxy proxy = null;
        if(StrUtil.isNotEmpty(chatGPTProperties.getProxyUrl()) && chatGPTProperties.getProxyPort() != 0){
            proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(chatGPTProperties.getProxyUrl(), chatGPTProperties.getProxyPort()));
        }
        HttpLoggingInterceptor httpLoggingInterceptor = new HttpLoggingInterceptor(new OpenAILogger());
        //!!!!!!测试或者发布到服务器千万不要配置Level == BODY!!!!
        //!!!!!!测试或者发布到服务器千万不要配置Level == BODY!!!!
        httpLoggingInterceptor.setLevel(HttpLoggingInterceptor.Level.HEADERS);
        OkHttpClient.Builder builder = new OkHttpClient.Builder()
                .addInterceptor(httpLoggingInterceptor)
                .connectTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(600, TimeUnit.SECONDS)
                .readTimeout(600, TimeUnit.SECONDS);
        if(null != proxy){
            builder.proxy(proxy);
        }
        OkHttpClient okHttpClient = builder.build();
        return OpenAiStreamClient
                .builder()
                .apiHost(chatGPTProperties.getApiHost())
                .apiKey(chatGPTProperties.getApiKeys())
                //自定义key使用策略 默认随机策略
                .keyStrategy(new KeyRandomStrategy())
                .okHttpClient(okHttpClient)
                .authInterceptor(new DynamicKeyOpenAiAuthInterceptor())
                .build();
    }
}
